# deep_learning_model.py
import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Dropout
from tensorflow.keras.callbacks import EarlyStopping
from sklearn.metrics import r2_score
import os
import matplotlib.pyplot as plt
import warnings
warnings.filterwarnings("ignore")

def build_dnn_model(input_shape):
    """
    构建一个用于回归任务的 DNN 模型
    :param input_shape: 输入维度
    :return: 编译好的 Keras 模型
    """
    model = Sequential([
        tf.keras.Input(shape=(input_shape,)),
        Dense(128, activation='relu', input_shape=(input_shape,)),
        Dropout(0.3),
        Dense(64, activation='relu'),
        Dense(1)  # 输出层（回归任务）
    ])

    model.compile(optimizer='adam',
                  loss='mean_squared_error',
                  metrics=['mae'])
    return model


def train_kfold_model(X, y, n_splits=5, epochs=50, batch_size=32, random_state=42):
    """
    使用 KFold 交叉验证训练 DNN 模型
    :param X: 特征数据
    :param y: 标签数据
    :param n_splits: 折数
    :param epochs: 训练轮数
    :param batch_size: 批量大小
    :param random_state: 随机种子
    :return: 最佳模型及评估指标
    """
    from sklearn.model_selection import KFold
    import numpy as np
    import joblib

    kf = KFold(n_splits=n_splits, shuffle=True, random_state=random_state)
    fold_no = 1
    val_scores = []
    mae_scores = []
    r2_scores = []

    for train_index, val_index in kf.split(X):
        print(f"\n 开始第 {fold_no} 折交叉验证")
        X_train, X_val = X[train_index], X[val_index]
        y_train, y_val = y[train_index], y[val_index]

        model = build_dnn_model(X.shape[1])

        early_stop = EarlyStopping(monitor='val_loss', patience=10, restore_best_weights=True)

        history = model.fit(X_train, y_train,
                            validation_data=(X_val, y_val),
                            epochs=epochs,
                            batch_size=batch_size,
                            callbacks=[early_stop],
                            verbose=0)

        y_pred = model.predict(X_val).flatten()
        mse = tf.keras.metrics.mean_squared_error(y_val, y_pred).numpy().mean()
        mae = tf.keras.metrics.mean_absolute_error(y_val, y_pred).numpy()
        r2 = r2_score(y_val, y_pred)

        val_scores.append(mse)
        mae_scores.append(mae)
        r2_scores.append(r2)

        print(f" 第 {fold_no} 折验证完成 - Val MSE: {mse:.6f}, MAE: {mae:.6f}, R²: {r2:.6f}")

        if fold_no == n_splits:
            best_model = model

        fold_no += 1

    avg_mse = np.mean(val_scores)
    avg_mae = np.mean(mae_scores)
    avg_r2 = np.mean(r2_scores)

    print("\n 五折平均评估指标：")
    print(f" 平均 MSE: {avg_mse:.6f}")
    print(f" 平均 MAE: {avg_mae:.6f}")
    print(f" 平均 R² Score: {avg_r2:.6f}")

    # 保存最佳模型
    output_dir = "../data/models"
    os.makedirs(output_dir, exist_ok=True)
    model_path = os.path.join(output_dir, "best_tensorflow_model.h5")
    best_model.save(model_path)
    print(f" 已保存最佳 TensorFlow 模型到 {model_path}")

    # 可视化 loss 曲线
    plot_history(history.history, output_dir)

    return best_model, avg_mse, avg_mae, avg_r2


def plot_history(history_dict, output_dir):
    """
    绘制训练过程中的 loss 曲线
    """
    plt.figure(figsize=(10, 5))
    plt.plot(history_dict['loss'], label='train_loss')
    plt.plot(history_dict['val_loss'], label='val_loss')
    plt.title('Model Loss Curve')
    plt.xlabel('Epochs')
    plt.ylabel('Loss')
    plt.legend()
    plt.grid(True)
    plt.savefig(os.path.join(output_dir, "tensorflow_training_history.png"))
    plt.close()


if __name__ == "__main__":
    import pandas as pd
    from sklearn.model_selection import train_test_split
    import numpy as np

    print(" 正在加载数据...")
    df = pd.read_parquet("../data/features/feature_engineered_data.parquet")

    # 检查并处理非数值列
    non_numeric_cols = df.select_dtypes(exclude='number').columns.tolist()
    if non_numeric_cols:
        print("发现非数值列，请先处理这些列：", non_numeric_cols)
        df = df.drop(columns=non_numeric_cols)

    X = df.drop(columns=['Average_Fare']).values.astype(np.float32)
    y = df['Average_Fare'].values.astype(np.float32)

    print(f" 数据形状：X={X.shape}, y={y.shape}")

    train_kfold_model(X, y)
